/*
 * transport.c 
 *
 * CS244a HW#3 (Reliable Transport)
 *
 * This file implements the STCP layer that sits between the
 * mysocket and network layers. You are required to fill in the STCP
 * functionality in this file. 
 *
 */


#include <stdio.h>
#include <stdarg.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>
#include <time.h>
#include <sys/time.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <signal.h>

#include "mysock.h"
#include "stcp_api.h"
#include "transport.h"


/* Default data offset (with no options in the header) */
#define STCP_DEFAULT_DATA_OFFSET    5
#define STCP_MAX_PACKET_LEN         (STCP_MSS + 16*4)   /*safer to be larger*/
#define STCP_RECEIVER_WINDOW_SIZE   3072
#define STCP_CONGESTION_WINDOW_SIZE 3072
#define STCP_MAX_SEQ_NUM            0xffffffff  /* maximum seq_num: 2^32-1 */
#define STCP_MAX_SEGMENT_LIFE       5   /* MSL 2 minutes */
#define STCP_RETRANSMISSION_LIMIT   5   /* number of retransmissions allowed */
#define STCP_INITIAL_RTT            1000 /* in milliseconds */
#define STCP_TIMEOUT_ALPHA          0.85
#define STCP_MIN_TIMEOUT            100
#define STCP_MAX_TIMEOUT            (STCP_MAX_SEGMENT_LIFE * 2 * 1000)
#define STCP_MIN_PAYLOAD_LEN        100


#define TCP_HEADER_LEN(p) (((STCPHeader *) p)->th_off * sizeof(uint32_t))

#define TCP_SEQ_NUM(p)    (ntohl(((STCPHeader *)p)->th_seq))
#define TCP_PAYLOAD_LEN(p, total_len)   (total_len - TCP_HEADER_LEN(p))

/* defines where the new packet is in the receiving buffer */
#define SEQ_NOT_IN          0
#define SEQ_HEAD_IN         1
#define SEQ_TAIL_IN         2
#define SEQ_BOTH_IN         (SEQ_HEAD_IN | SEQ_TAIL_IN)
#define SEQ_EMPTY_PACKET    4

#define ACK_ACCEPTABLE      0
#define ACK_EMPTY           1
#define ACK_UNACCEPTABLE    2

#define PACKET_SENT         1
#define PACKET_RCVD         2

#define EVENT_BLOCK_APP_DATA (ANY_EVENT & (~APP_DATA))

#define STCP_NO_ERROR       0
#define STCP_ERROR          1

/*
#define _DEBUG_PRINT
*/

/* connection states */
enum 
{ 
    CLOSED,
    LISTEN,
    SYN_RCVD,
    SYN_SENT,
    ESTABLISHED,
    FIN_WAIT1,
    FIN_WAIT2,
    CLOSE_WAIT,
    CLOSING,
    LAST_ACK,
    TIME_WAIT
};    

typedef struct buffer_queue_node
{
    tcp_seq  seq;   /* starting sequence number */
    int      len;
    uint8_t *data;  /* the real data packet */
    buffer_queue_node *next;
} buffer_queue_node_t;

typedef struct 
{
    buffer_queue_node_t *head;
} buffer_queue_t;

typedef struct timer_queue_node
{
    uint8_t          *packet;     /* header included */
    int               len;        
    struct timespec   timer_start;    /* the time the timer started */
    pthread_t         timer_thread;
    pthread_cond_t    timer_kill_cond;
    pthread_mutex_t   timer_kill_lock;
    bool_t            timer_killed;
    int               num_trials_left;
    timer_queue_node *next;
} timer_queue_node_t;

typedef struct
{
    timer_queue_node_t *head;
    timer_queue_node_t *tail;
} timer_queue_t;

/* this structure is global to a mysocket descriptor */
typedef struct
{
    bool_t done;    /* TRUE once connection is closed */

    int connection_state;   /* state of the connection (established, etc.) */

    tcp_seq initial_sequence_num;   /* send ISS */
    tcp_seq send_unack;
    tcp_seq send_next;
    tcp_seq recv_next;
    tcp_seq recv_initial_seq_num;

    unsigned int recv_window_size;
    unsigned int send_window_size;

    uint8_t      send_buffer[STCP_MAX_PACKET_LEN];
    unsigned int send_buf_next;
    unsigned int event_mask;

    unsigned int rtt;
    unsigned int timeout;

    buffer_queue_t *recv_buffer;
    timer_queue_t  *retransmission_queue;

    /* only one lock, so deadlock is impossible */
    pthread_mutex_t retransmission_lock;

    bool_t error;

    /* any other connection-wide global variables go here */
} context_t;

typedef struct 
{
    context_t   *ctx;
    mysocket_t   sd;
    unsigned int timeout_msec;
    timer_queue_node_t *node;
} timer_arg_type;



/* function declarations */

void print_buffer_queue(context_t *ctx);
void print_tcphdr_byte(uint8_t *tcphdr);
void print_tcphdr(STCPHeader *tcphdr);
void print_packet_info(void const *packet, unsigned int recv_or_send);

void fill_in_tcphdr(STCPHeader *tcphdr, tcp_seq seq, tcp_seq ack, 
        uint8_t offset, uint8_t flags, uint16_t window);

unsigned int ack_acceptable(context_t *ctx, tcp_seq ack);
unsigned int seq_acceptable(context_t *ctx, tcp_seq seq, int segmentlen);

void   clean_send_buffer(mysocket_t sd, context_t *ctx);
ssize_t send_flag_packet(mysocket_t sd, context_t *ctx, uint8_t flags);
void process_rcvd_packet(mysocket_t sd, context_t *ctx, 
        uint8_t *packet, int packet_len);
void insert_packet_to_buffer(context_t *ctx, tcp_seq seq_start, int seg_len, 
        uint8_t *packet);
int get_first_packet_from_buffer(context_t *ctx, uint8_t *buf, int maxlen);

ssize_t send_and_start_timer(mysocket_t sd, context_t *ctx, const void *tcphdr,
        size_t hdr_len, const void *payload, size_t payload_len);
/* void *timer_thread_func(void *arg); */
void notify_retransmission_queue(context_t *ctx, tcp_seq ack);
void cleanup_timer_node(timer_queue_node_t *node);
void start_timer(mysocket_t sd, context_t *ctx, timer_queue_node_t *node);

static void generate_initial_seq_num(context_t *ctx);
static void control_loop(mysocket_t sd, context_t *ctx);


/* initialize context parameters, it is initially zeroed out */
void context_init(context_t *ctx)
{
    memset(ctx, 0, sizeof(context_t));

    ctx->done = FALSE;
    ctx->connection_state = CLOSED;

    generate_initial_seq_num(ctx);
    ctx->send_unack = ctx->initial_sequence_num;
    ctx->send_next = ctx->initial_sequence_num;
    ctx->recv_next = 0;
    ctx->recv_initial_seq_num = 0;
    ctx->recv_window_size = STCP_RECEIVER_WINDOW_SIZE;
    ctx->send_window_size = STCP_CONGESTION_WINDOW_SIZE;

    ctx->send_buf_next = 0;

    ctx->event_mask = ANY_EVENT;

    ctx->rtt = STCP_INITIAL_RTT;   
    ctx->timeout = MAX(2 * ctx->rtt, STCP_MIN_TIMEOUT);

    ctx->recv_buffer = (buffer_queue_t *)calloc(1, sizeof(buffer_queue_t));
    ctx->retransmission_queue = (timer_queue_t *)
        calloc(1, sizeof(timer_queue_t));

    pthread_mutex_init(&ctx->retransmission_lock, NULL);

    ctx->error = STCP_NO_ERROR;
}

/* destroy resources associated with this context */
void context_destroy(context_t *ctx)
{
    buffer_queue_t *bq = ctx->recv_buffer;
    buffer_queue_node_t *p;

    timer_queue_t *rq = ctx->retransmission_queue;
    timer_queue_node_t *q;

    /* destroy the receive buffer */

#ifdef _DEBUG_PRINT
    printf("Destroying the receive buffer...\n");
#endif

    while (bq->head)
    {
        p = bq->head;
        bq->head = p->next;

        free(p->data);
        free(p);
    }

    free(ctx->recv_buffer);

    /* destroy the retransmission queue */

#ifdef _DEBUG_PRINT
    printf("Destroying the retransmission queue...\n");
#endif


    pthread_mutex_lock(&ctx->retransmission_lock);
    while (rq->head)
    {
        q = rq->head;
        rq->head = q->next;
        pthread_mutex_unlock(&ctx->retransmission_lock);

        cleanup_timer_node(q);
        free(q);

        pthread_mutex_lock(&ctx->retransmission_lock);
    }
    pthread_mutex_unlock(&ctx->retransmission_lock);

    free(ctx->retransmission_queue);

    pthread_mutex_destroy(&ctx->retransmission_lock);
}

/* initialise the transport layer, and start the main loop, handling
 * any data from the peer or the application.  this function should not
 * return until the connection is closed.
 */
void transport_init(mysocket_t sd, bool_t is_active)
{
    context_t *ctx = (context_t *) calloc(1, sizeof(context_t));
    assert(ctx);

#ifdef _DEBUG_PRINT
    printf("***** Initializing STCP *****\n");
#endif

    context_init(ctx);

#ifdef _DEBUG_PRINT
    printf("Initial seq num: %d\n", ctx->initial_sequence_num);
#endif

    /* XXX: you should send a SYN packet here if is_active, or wait for one
     * to arrive if !is_active.  after the handshake completes, unblock the
     * application with stcp_unblock_application(sd).  you may also use
     * this to communicate an error condition back to the application, e.g.
     * if connection fails; to do so, just set errno appropriately (e.g. to
     * ECONNREFUSED, etc.) before calling the function.
     */

    if (is_active)
    {
        /* send SYN packet */
        send_flag_packet(sd, ctx, TH_SYN);

        /* SYN takes up 1 byte of sequence space */
        ctx->send_next++;
        ctx->connection_state = SYN_SENT;
    }
    else
    {
        ctx->connection_state = LISTEN;
    }

    control_loop(sd, ctx);

#ifdef _DEBUG_PRINT
    printf("Connection for sd %d closed...\n", sd);
#endif

    /* do any cleanup here */
    context_destroy(ctx);
    free(ctx);

    if (ctx->error == STCP_ERROR)
    {
        if (ctx->connection_state == LISTEN || 
                ctx->connection_state == SYN_SENT ||
                ctx->connection_state == SYN_RCVD)
        {
            errno = ECONNREFUSED;
            stcp_unblock_application(sd);
        }
        else
        {
            errno = ECONNRESET;
        }
    }
}

/* check if the ACK is acceptable */
unsigned int ack_acceptable(context_t *ctx, tcp_seq ack)
{
    if (ctx->send_unack == ack)
        return ACK_EMPTY; 
    if (ctx->send_unack <= ctx->send_next)
        return (ctx->send_unack < ack && ack <= ctx->send_next ?
                ACK_ACCEPTABLE : ACK_UNACCEPTABLE);
    else
        return (ctx->send_unack < ack || ack <= ctx->send_next ?
                ACK_ACCEPTABLE : ACK_UNACCEPTABLE);
}

/* check if the sequence number is acceptable, return sequence state,
 * head in, tail in, etc. */
unsigned int seq_acceptable(context_t *ctx, tcp_seq seq, int segmentlen)
{
    bool_t head_in, tail_in;
    unsigned int rc = 0;

    if (STCP_MAX_SEQ_NUM - ctx->recv_next < ctx->recv_window_size)
    {
        head_in = (ctx->recv_next <= seq || 
                seq < ctx->recv_next + ctx->recv_window_size);
        tail_in = (ctx->recv_next <= seq + segmentlen - 1 || 
                seq + segmentlen - 1 < ctx->recv_next + ctx->recv_window_size);
    }
    else
    {
        head_in = (ctx->recv_next <= seq &&
                seq < ctx->recv_next + ctx->recv_window_size);
        tail_in = (ctx->recv_next <= seq + segmentlen - 1 && 
                seq + segmentlen - 1 < ctx->recv_next + ctx->recv_window_size);
    }

    if (segmentlen == 0)
        if ((ctx->recv_window_size == 0 && ctx->recv_next == seq) ||
                (ctx->recv_window_size > 0 && head_in))
            return SEQ_EMPTY_PACKET;
        else
            return SEQ_NOT_IN;
    else
        if (ctx->recv_window_size == 0)
            return SEQ_NOT_IN;
        else
        {
            if (head_in)
                rc |= SEQ_HEAD_IN;
            if (tail_in)
                rc |= SEQ_TAIL_IN;
            return rc;
        }
}

/* fill in STCP header with required fields */
void fill_in_tcphdr(STCPHeader *tcphdr, 
        tcp_seq seq, 
        tcp_seq ack, 
        uint8_t offset,
        uint8_t flags, 
        uint16_t window)
{
    memset(tcphdr, 0, sizeof(STCPHeader));
    tcphdr->th_seq = htonl(seq);
    tcphdr->th_ack = htonl(ack);
    tcphdr->th_off = offset;
    tcphdr->th_flags = flags;
    tcphdr->th_win = htons(window);

}

/* generate random initial sequence number for an STCP connection */
static void generate_initial_seq_num(context_t *ctx)
{
    assert(ctx);

    struct timeval tp = { 0, 0 };
    gettimeofday(&tp, NULL);

    /* time(NULL) is not accurate enough to make the initial seq_num 
     * different, so use gettimeofday instead, which can get the time in
     * microseconds.
     */
    srand(tp.tv_usec);  

#ifdef FIXED_INITNUM
    /* please don't change this! */
    ctx->initial_sequence_num = 1;
#else
    /* you have to fill this up */
    ctx->initial_sequence_num = rand() & 0xff;
#endif
}

void print_state(context_t *ctx)
{
    assert(ctx);

    switch(ctx->connection_state)
    {
        case CLOSED:        printf("[CLOSED]\n"); break;
        case LISTEN:        printf("[LISTEN]\n"); break;
        case SYN_RCVD:      printf("[SYN_RCVD]\n"); break;
        case SYN_SENT:      printf("[SYN_SENT]\n"); break;
        case ESTABLISHED:   printf("[ESTABLISHED]\n"); break;
        case FIN_WAIT1:     printf("[FIN_WAIT1]\n"); break;
        case FIN_WAIT2:     printf("[FIN_WAIT2]\n"); break;
        case CLOSE_WAIT:    printf("[CLOSE_WAIT]\n"); break;
        case CLOSING:       printf("[CLOSING]\n"); break;
        case LAST_ACK:      printf("[LAST_ACK]\n"); break;
        case TIME_WAIT:     printf("[TIME_WAIT]\n"); break;
        default:            printf("[ERROR_STATE]\n"); break;
    }
}

void print_packet_info(const void *packet, unsigned int recv_or_send)
{

    STCPHeader *tcphdr = (STCPHeader *)packet;

    switch (recv_or_send)
    {
        case PACKET_SENT:
            printf("<<<< SENT, ");
            break;
        case PACKET_RCVD:
            printf(">>>> RCVD, ");
            break;
        default:
            printf("!!!! Error recv_or_send!\n");
            return;
    }
    
   
    printf("seq %d, ack %d, flags [", ntohl(tcphdr->th_seq), 
            ntohl(tcphdr->th_ack));

    if (tcphdr->th_flags & TH_FIN)
        printf(" FIN");
    if (tcphdr->th_flags & TH_SYN)
        printf(" SYN");
    if (tcphdr->th_flags & TH_RST)
        printf(" RST");
    if (tcphdr->th_flags & TH_PUSH)
        printf(" PUSH");
    if (tcphdr->th_flags & TH_ACK)
        printf(" ACK");
    if (tcphdr->th_flags & TH_URG)
        printf(" URG");

    printf(" ], window %d\n", ntohs(tcphdr->th_win));
}

/* control_loop() is the main STCP loop; it repeatedly waits for one of the
 * following to happen:
 *   - incoming data from the peer
 *   - new data from the application (via mywrite())
 *   - the socket to be closed (via myclose())
 *   - a timeout
 */
static void control_loop(mysocket_t sd, context_t *ctx)
{
    assert(ctx);

    uint8_t buf[STCP_MAX_PACKET_LEN];
    ssize_t bytes_read;

    struct timespec abstime;
    struct timeval _temptime;
    
    unsigned int event = ANY_EVENT;
    size_t available_buf_size;
    size_t bytes_sent_in_window;

#ifdef _DEBUG_PRINT
    printf("In the control loop!\n");
#endif

    while (!ctx->done)
    {
        /* don't print for probing events for ctx->done */
#ifdef _DEBUG_PRINT
        if (event != TIMEOUT)
            print_state(ctx);
#endif

        /* I should change makefile to make clock_gettime work, this is a
         * workaround */
        gettimeofday(&_temptime, NULL);
        abstime.tv_sec = _temptime.tv_sec;
        abstime.tv_nsec = _temptime.tv_usec * 1000;

        if (ctx->connection_state == TIME_WAIT)
        {
            abstime.tv_sec += STCP_MAX_SEGMENT_LIFE * 2;
#ifdef _DEBUG_PRINT
            printf("Time: %s\n", asctime(gmtime(&(_temptime.tv_sec))));
#endif
            event = stcp_wait_for_event(
                    sd, ANY_EVENT & ctx->event_mask, &abstime);
        }
        else
        {
            /* set timeout as well, to make sure the process is not stuck so
             * it can respond to events like retransmission failure */
            abstime.tv_sec += 1;
            event = stcp_wait_for_event(
                    sd, ANY_EVENT & ctx->event_mask, &abstime);
        }

        /* data received from application, send it directly if there is space
         * in send buffer */
        if (event & APP_DATA)
        {
            /* number of bytes already sent/buffered in the window */
            bytes_sent_in_window = ctx->send_next - ctx->send_unack +
                ctx->send_buf_next;

            /* no space in the buffer, block the application */
            if (bytes_sent_in_window >= ctx->send_window_size)
            {
#ifdef _DEBUG_PRINT
                printf("Application blocked...\n");
#endif
                ctx->event_mask = EVENT_BLOCK_APP_DATA;
            }
            else
            {
                /* this may happen when the receiver shrinked its window */
                if (ctx->send_window_size < bytes_sent_in_window)
                    available_buf_size = 0;
                else
                    available_buf_size = MIN(
                            ctx->send_window_size - bytes_sent_in_window, 
                            STCP_MSS - ctx->send_buf_next);

                bytes_read = stcp_app_recv(sd, 
                        ctx->send_buffer + ctx->send_buf_next,
                        available_buf_size);

                ctx->send_buf_next += bytes_read;

                if (ctx->send_buf_next == STCP_MSS ||
                        ctx->send_unack == ctx->send_next)
                    clean_send_buffer(sd, ctx);
                else if (bytes_read > 0)
#ifdef _DEBUG_PRINT
                    printf("[DATA FROM APP] %d bytes buffered, buf_len %u...\n",
                            bytes_read, ctx->send_buf_next);
#else
                {;}
#endif
            }
        }
        
        /* data received from the peer, send it to the application */
        if (event & NETWORK_DATA)
        {
            /* XXX: problem in _network_get_interface_ip, use cdf machines and
             * it works fine. */

            /* header included in the data as well */
            bytes_read = stcp_network_recv(sd, buf, STCP_MAX_PACKET_LEN);

#ifdef _DEBUG_PRINT
            print_packet_info(buf, PACKET_RCVD);
#endif

            process_rcvd_packet(sd, ctx, buf, bytes_read);
        }

        if (event & APP_CLOSE_REQUESTED)
        {
#ifdef _DEBUG_PRINT
            printf("[CLOSE_REQUESTED]\n");
#endif

            /* close requested before connection established */
            if (ctx->connection_state == LISTEN || 
                    ctx->connection_state == SYN_SENT)
            {
                ctx->connection_state = CLOSED;
                ctx->done = TRUE;
                break;
            }

            send_flag_packet(sd, ctx, TH_FIN | TH_ACK);

            ctx->send_next++;

#ifdef _DEBUG_PRINT
            printf("[FIN SENT]\n");
#endif

            if (ctx->connection_state == ESTABLISHED || 
                    ctx->connection_state == SYN_RCVD)
                ctx->connection_state = FIN_WAIT1;
            else if (ctx->connection_state == CLOSE_WAIT)
                ctx->connection_state = LAST_ACK;
            /* other connection states are impossible */
        }

        if (event == TIMEOUT && ctx->connection_state == TIME_WAIT)
        {
#ifdef _DEBUG_PRINT
            printf("[TIMEOUT]\n");
#endif
            ctx->connection_state = CLOSED;
#ifdef _DEBUG_PRINT
            printf("[CONNECTION CLOSED]\n");
#endif
            ctx->done = TRUE;
        }
    }
}



/* This functio processes packet received from the network, it chooses the
 * right operation based on current state of the connection. */
void process_rcvd_packet(mysocket_t sd, context_t *ctx, 
        uint8_t *packet, int packet_len)
{
    STCPHeader *tcphdr = (STCPHeader *)packet;

    tcp_seq cur_seq = ntohl(tcphdr->th_seq);
    tcp_seq cur_ack = ntohl(tcphdr->th_ack);

    int seg_len = packet_len - TCP_HEADER_LEN(tcphdr);

    tcp_seq segment_effective_seq_start = 0;
    int     segment_effective_len       = 0;

    unsigned int seq_where = seq_acceptable(ctx, cur_seq, seg_len);
    unsigned int ack_state = ack_acceptable(ctx, cur_ack);

    uint8_t buf[STCP_RECEIVER_WINDOW_SIZE];
    int     bytes_read;

    /* Indicating whether all waiting data is acknowledged, used to process
     * FIN packets */
    bool_t ack_cleared = (cur_ack == ctx->send_next);

#ifdef _DEBUG_PRINT
    switch(ack_state)
    {
        case ACK_UNACCEPTABLE:  printf("ack unacceptable"); break;
        case ACK_EMPTY:         printf("ack empty");        break;
        case ACK_ACCEPTABLE:    printf("ack acceptable");   break;
        default: printf("ack check error");
    }
    printf(", seq");
    if (seq_where & SEQ_HEAD_IN)
        printf(" head_in");
    if (seq_where & SEQ_TAIL_IN)
        printf(" tail_in");
    if (seq_where & SEQ_EMPTY_PACKET)
        printf(" empty_packet");
    if (seq_where == SEQ_NOT_IN)
        printf(" not in");
    printf("\n");

    if (ack_cleared)
        printf("[ACK CLEARED]\n");
#endif

    switch(ctx->connection_state)
    {
        case CLOSED: 
            break;
        case LISTEN: 
            if (tcphdr->th_flags & TH_SYN)
            {
                ctx->recv_next = ntohl(tcphdr->th_seq) + 1;
                ctx->recv_initial_seq_num = ntohl(tcphdr->th_seq);
                ctx->send_window_size = MIN(STCP_CONGESTION_WINDOW_SIZE,
                        ntohs(tcphdr->th_win));

                send_flag_packet(sd, ctx, TH_SYN | TH_ACK);

                ctx->send_next++;
                ctx->connection_state = SYN_RCVD;
            }
            break;
        case SYN_SENT: 
            if ((!(tcphdr->th_flags & TH_ACK) ||
                    (ack_state != ACK_UNACCEPTABLE)) &&
                    tcphdr->th_flags & TH_SYN)
            {
                ctx->recv_next = cur_seq + 1;
                ctx->recv_initial_seq_num = cur_seq;
                ctx->send_window_size = MIN(STCP_CONGESTION_WINDOW_SIZE,
                        ntohs(tcphdr->th_win));

                if (tcphdr->th_flags & TH_ACK)
                {
                    ctx->send_unack = cur_ack;
                    notify_retransmission_queue(ctx, cur_ack);
                    clean_send_buffer(sd, ctx);
                    ctx->event_mask = ANY_EVENT;
                }

                if (ctx->initial_sequence_num < ctx->send_unack)
                {
                    ctx->connection_state = ESTABLISHED;
                    send_flag_packet(sd, ctx, TH_ACK);

                    stcp_unblock_application(sd);
                }
                else
                {
                    ctx->connection_state = SYN_RCVD;
                    send_flag_packet(sd, ctx, TH_SYN | TH_ACK);
                }
            }
            break;
        default:
            /* check sequence number, send ACK and drop it if unacceptable */
            if (seq_where == SEQ_NOT_IN)
            {
#ifdef _DEBUG_PRINT
                printf("Bad sequence number: %d...\n", cur_seq);
#endif
                send_flag_packet(sd, ctx, TH_ACK);
                return;
            }
            /* check ACK bit, if off, drop the packet */
            if (!(tcphdr->th_flags & TH_ACK))
            {
#ifdef _DEBUG_PRINT
                printf("ACK bit not set, packet dropped...\n");
#endif
                return;
            }

            /* process the ACK bit */
            switch (ctx->connection_state)
            {
                case SYN_RCVD:
                    if (ack_state != ACK_UNACCEPTABLE)
                    {
                        ctx->connection_state = ESTABLISHED;
                        stcp_unblock_application(sd);
                    }
                    else
                        return;

                case ESTABLISHED:
                case FIN_WAIT1:
                case FIN_WAIT2:
                case CLOSE_WAIT:
                case CLOSING:
                    if (ack_state == ACK_UNACCEPTABLE)
                    {
#ifdef _DEBUG_PRINT
                        printf("Duplicate/Bad ACK...\n");
#endif
                        send_flag_packet(sd, ctx, TH_ACK);
                        return;
                    }

                    ctx->send_unack = cur_ack;
                    ctx->send_window_size = MIN(STCP_CONGESTION_WINDOW_SIZE,
                            ntohs(tcphdr->th_win));

                    /* valid ack, clean the send buffer */
                    clean_send_buffer(sd, ctx);
                    ctx->event_mask = ANY_EVENT;

                    if (ack_state == ACK_ACCEPTABLE)
                        notify_retransmission_queue(ctx, cur_ack);
                    
                    /* state transitions */

                    /* FIN acknowledged */
                    if (ctx->connection_state == FIN_WAIT1 && ack_cleared)
                        ctx->connection_state = FIN_WAIT2;

                    /* user's transmission finished */
                    if (ctx->connection_state == FIN_WAIT2) {;}

                    if (ctx->connection_state == CLOSING && ack_cleared)
                        ctx->connection_state = TIME_WAIT;
                    break;
                case LAST_ACK:
                    /* transmission closed on both sides */
                    ctx->send_unack = cur_ack;
                    if (ack_cleared)
                    {
                        ctx->connection_state = CLOSED;
                        ctx->done = TRUE;
                        return;
                    }
                    break;
                case TIME_WAIT:
                    /* 2 MSL timer will be started in control_loop */
                    send_flag_packet(sd, ctx, TH_ACK);
                    return;
                default: 
                    break;
            }

            /* process the segment */
            switch (ctx->connection_state)
            {
                case ESTABLISHED:
                case FIN_WAIT1:
                case FIN_WAIT2:
                    /* put the segment into the receive buffer */

                    if (seq_where == SEQ_EMPTY_PACKET)
                    {
                        if (ack_state != ACK_EMPTY) /* avoid loop */
                            send_flag_packet(sd, ctx, TH_ACK);
                        break;
                    }

                    /* packets not empty should be put in the buffer */
                    if (seq_where & SEQ_HEAD_IN)
                        segment_effective_seq_start = cur_seq;
                    else
                        segment_effective_seq_start = ctx->recv_next;

                    if (seq_where & SEQ_TAIL_IN)
                        segment_effective_len = cur_seq + seg_len - 
                            segment_effective_seq_start;
                    else
                        segment_effective_len = ctx->recv_window_size - 
                            (segment_effective_seq_start - ctx->recv_next);

                    insert_packet_to_buffer(ctx, segment_effective_seq_start,
                            segment_effective_len, 
                            packet + TCP_HEADER_LEN(tcphdr) + 
                            segment_effective_seq_start - cur_seq);

                    /* if the effective block is the block we are waiting for,
                     * deliver it to the application */
                    if (segment_effective_seq_start == ctx->recv_next)
                    {
                        bytes_read = get_first_packet_from_buffer(ctx, buf,
                                STCP_RECEIVER_WINDOW_SIZE);
                        if (bytes_read > 0)
                        {
                            ctx->recv_next += bytes_read;
                            stcp_app_send(sd, buf, bytes_read);

#ifdef _DEBUG_PRINT
                            printf("%d bytes send to app...\n", bytes_read);
#endif
                        }
                    }

                    send_flag_packet(sd, ctx, TH_ACK);

                    break;
                default:
                    break;
            }

            /* check the FIN bit */
            if (!(tcphdr->th_flags & TH_FIN))
                return;

            /* tell the application FIN received */
            stcp_fin_received(sd);
#ifdef _DEBUG_PRINT
            printf("[FIN RECEIVED]\n");
#endif

            /* FIN takes up 1 byte of seq space */
            ctx->recv_next++;   

            /* this will be different from the ACK sent when processing the
             * segment, because recv_next is different */
            send_flag_packet(sd, ctx, TH_ACK);  

            switch(ctx->connection_state)
            {
                case SYN_RCVD:
                case ESTABLISHED:
                    ctx->connection_state = CLOSE_WAIT;
                    break;
                case FIN_WAIT1:
                    if (ack_cleared)
                    {
                        ctx->connection_state = TIME_WAIT;
                        /* 2MSL timer will be started in process_rcvd_packet */
                    }
                    else
                        ctx->connection_state = CLOSING;
                    break;
                case FIN_WAIT2:
                    ctx->connection_state = TIME_WAIT;
                    break;
                case TIME_WAIT:
                    break;
                default:
                    break;
            }

            break;
    }
}

/* send buffered data in the send buffer */
void clean_send_buffer(mysocket_t sd, context_t *ctx)
{
    STCPHeader _tcphdr;
    STCPHeader *tcphdr = &_tcphdr;
    ssize_t bytes_sent;
    ssize_t bytes_to_send;
    uint8_t _buf[STCP_MAX_PACKET_LEN];
    
    /* window shrinked, can't send any data now */
    if (ctx->send_window_size < ctx->send_next - ctx->send_unack)
        return;

    /* no data to send */
    if (ctx->send_buf_next == 0)
        return;

    /* bytes to send is the minimum of the current buffer length and the
     * available window space */
    bytes_to_send = MIN(ctx->send_buf_next,
            ctx->send_window_size - (ctx->send_next - ctx->send_unack));

    fill_in_tcphdr(tcphdr, ctx->send_next, ctx->recv_next, 
            STCP_DEFAULT_DATA_OFFSET, TH_ACK, ctx->recv_window_size);

    bytes_sent = send_and_start_timer(sd, ctx, tcphdr, 
            sizeof(STCPHeader), ctx->send_buffer, bytes_to_send);

#ifdef _DEBUG_PRINT
    printf("[DATA FROM APP] %d bytes sent to peer...\n", bytes_sent);
#endif

    assert(bytes_sent >= 0 && (unsigned)bytes_sent == sizeof(STCPHeader) + 
            bytes_to_send);

    memcpy(_buf, ctx->send_buffer + bytes_to_send,
            ctx->send_buf_next - bytes_to_send);
    memcpy(ctx->send_buffer, _buf, 
            ctx->send_buf_next - bytes_to_send);
    
    ctx->send_next += bytes_to_send;
    ctx->send_buf_next -= bytes_to_send;
}

/* get the first packet from buffer queue, return the number of bytes read */
int get_first_packet_from_buffer(context_t *ctx, uint8_t *buf, int maxlen)
{
    int bytes_read;
    
    buffer_queue_t *bq = ctx->recv_buffer;
    buffer_queue_node_t *q;

#ifdef _DEBUG_PRINT
    printf("get_first_packet_from_buffer()\n");
#endif

    if (bq->head && bq->head->len > maxlen)
    {
        memcpy(buf, bq->head->data, maxlen);
        memcpy(bq->head->data, bq->head->data + maxlen, 
                bq->head->len - maxlen);
        bytes_read = maxlen;
    }
    else if (bq->head)
    {
        memcpy(buf, bq->head->data, bq->head->len);
        bytes_read = bq->head->len;
        
        q = bq->head;
        bq->head = q->next;
        free(q->data);
        free(q);
    }
    else
    {
        bytes_read = 0;
    }
    return bytes_read;
}

/* insert a packet into the receiving buffer */
void insert_packet_to_buffer(context_t *ctx, tcp_seq seq_start, int seg_len, 
        uint8_t *packet)
{
    assert(ctx);

    buffer_queue_t *bq = ctx->recv_buffer;
    buffer_queue_node_t *node;

    buffer_queue_node_t *p_cur;
    buffer_queue_node_t *p_prev;
    buffer_queue_node_t *p_from;
    buffer_queue_node_t *p_to;

    int new_len = 0;

#ifdef _DEBUG_PRINT
    printf("insert_packet_to_buffer()\n");
#endif

    node = (buffer_queue_node_t *) calloc(1, sizeof(buffer_queue_node_t));
    node->seq = seq_start;
    node->len = seg_len;
    node->data = (uint8_t *) calloc(seg_len, sizeof(uint8_t));
    memcpy(node->data, packet, seg_len);

#ifdef _DEBUG_PRINT
    print_buffer_queue(ctx);
#endif

    /* insert the new node into the queue */
    p_cur = bq->head;
    p_prev = NULL;
    while (p_cur && p_cur->seq - ctx->recv_next < seq_start - ctx->recv_next)
    {
        p_prev = p_cur;
        p_cur = p_cur->next;
    }

    node->next = p_cur;
    if (p_prev)
        p_prev->next = node;
    else
        bq->head = node;

#ifdef _DEBUG_PRINT
    print_buffer_queue(ctx);
#endif

    /* merge nodes to form a larger block */
    p_cur = bq->head;
    while (p_cur->seq - ctx->recv_next + p_cur->len < 
            seq_start - ctx->recv_next)
        p_cur = p_cur->next;

    p_from = p_cur;

    while (p_cur && p_cur->seq - ctx->recv_next <=
            seq_start - ctx->recv_next + seg_len)
    {
        p_to = p_cur;
        p_cur = p_cur->next;
    }

    new_len = (p_to->seq - ctx->recv_next + p_to->len) - 
        (p_from->seq - ctx->recv_next);

    p_from->data = (uint8_t *)realloc(p_from->data, new_len);
    memcpy(p_from->data + (node->seq - p_from->seq), node->data, seg_len);
    memcpy(p_from->data + (node->seq - p_from->seq + seg_len), 
            p_to->data + (node->seq - ctx->recv_next + seg_len) - 
            (p_to->seq - ctx->recv_next), 
            new_len - (node->seq - p_from->seq + seg_len));
    p_from->len = new_len;

#ifdef _DEBUG_PRINT
    print_buffer_queue(ctx);
#endif

    /* remove redundant nodes */
    p_cur = p_from->next;
    while (p_cur && p_cur->seq - ctx->recv_next <= 
            p_from->seq - ctx->recv_next + new_len)
    {
        p_prev = p_cur;
        p_cur = p_cur->next;

        free(p_prev->data);
        free(p_prev);
        p_from->next = p_cur;
    }

#ifdef _DEBUG_PRINT
    print_buffer_queue(ctx);
#endif
}

/* only for debug */
void print_buffer_queue(context_t *ctx)
{
    assert(ctx);

    buffer_queue_node_t *p = ctx->recv_buffer->head;

    while (p)
    {
        printf("[seq:%d, len:%d]->", p->seq, p->len);
        p = p->next;
    }
    printf("NULL\n");
}

/* compute the length of the time from t_start to t_end, in milliseconds */
unsigned int diff_msec(struct timespec *t_start, struct timespec *t_end)
{
    unsigned int tdiff = 0;
    unsigned int borrow = 0;

    if (t_end->tv_nsec < t_start->tv_nsec)
        borrow = 1;
    tdiff = borrow * 1000 - t_start->tv_nsec / 1000000 + 
        t_end->tv_nsec / 1000000;

    if (t_end->tv_sec < t_start->tv_sec || 
            (t_end->tv_sec == t_start->tv_sec && borrow))
        return 0;
    
    tdiff += (t_end->tv_sec - t_start->tv_sec - borrow) * 1000;
    return tdiff;
}

/* perform cleanup after a timer and corresponding packet can be removed, the
 * calling thread should have ctx->retransmission_lock locked during the call */
void cleanup_timer_node(timer_queue_node_t *node)
{
    /* kill the timer thread */
    pthread_mutex_lock(&node->timer_kill_lock);
    node->timer_killed = TRUE;
    pthread_cond_broadcast(&node->timer_kill_cond);
    pthread_mutex_unlock(&node->timer_kill_lock);

    pthread_join(node->timer_thread, NULL);

    /* cleanup and free resources */
    free(node->packet);
    pthread_cond_destroy(&node->timer_kill_cond);
    pthread_mutex_destroy(&node->timer_kill_lock);
}

/* Remove packets that are ACK'ed by the given ACK number, reset RTT according
 * to the recorded time of the corresponding packet */
void notify_retransmission_queue(context_t *ctx, tcp_seq ack)
{
    timer_queue_t *rq;
    timer_queue_node_t *q;
    struct timeval  abstime;
    struct timespec cur_time;

#ifdef _DEBUG_PRINT
    printf("notifying retransmission queue...\n");
#endif

    rq = ctx->retransmission_queue;

    pthread_mutex_lock(&ctx->retransmission_lock);
    /* remove completely acknowledged packets */
    for (q = rq->head; q && TCP_SEQ_NUM(q->packet) + 
            TCP_PAYLOAD_LEN(q->packet, q->len) <= ack;)
    {
        /* update RTT if ACK arrived on time */
        if (TCP_SEQ_NUM(q->packet) + TCP_PAYLOAD_LEN(q->packet, q->len) == ack)
        {
            gettimeofday(&abstime, NULL);

            cur_time.tv_sec = abstime.tv_sec;
            cur_time.tv_nsec = abstime.tv_usec * 1000;

            /* only update when it is the first transmission */
            if (q->num_trials_left == STCP_RETRANSMISSION_LIMIT)
            {
#ifdef _DEBUG_PRINT
                printf("RTT changed from %d, ", ctx->rtt);
#endif

                ctx->rtt = (unsigned int) (STCP_TIMEOUT_ALPHA * ctx->rtt +
                    (1 - STCP_TIMEOUT_ALPHA) * 
                    diff_msec(&(q->timer_start), &cur_time));
                ctx->timeout = MAX(2 * ctx->rtt, STCP_MIN_TIMEOUT);

#ifdef _DEBUG_PRINT
                printf("to %d, estimated RTT %d\n", ctx->rtt,
                        diff_msec(&(q->timer_start), &cur_time));
#endif
            }
        }

        /* remove the packet */
#ifdef _DEBUG_PRINT
        printf("Packet seq %d removed from retransmission queue...\n", 
                TCP_SEQ_NUM(q->packet));
#endif

        rq->head = q->next;
        pthread_mutex_unlock(&ctx->retransmission_lock);

        cleanup_timer_node(q);

        pthread_mutex_lock(&ctx->retransmission_lock);
        q = rq->head;
    }

    /* empty queue */
    if (!(rq->head))
        rq->tail = NULL;

    pthread_mutex_unlock(&ctx->retransmission_lock);
}

void *new_timer_thread_func(void *arg)
{
    timer_arg_type *targ = (timer_arg_type *)arg;

    context_t            *ctx = targ->ctx;
    mysocket_t             sd = targ->sd;
    timer_queue_node_t  *node = targ->node;
    unsigned int timeout_msec = targ->timeout_msec;

    struct timespec wakeuptime;
    struct timeval abstime;

    while (node->num_trials_left >= 0)
    {
        wakeuptime.tv_sec = timeout_msec / 1000;
        wakeuptime.tv_nsec = timeout_msec % 1000 * 1000000;

        gettimeofday(&abstime, NULL);

        wakeuptime.tv_sec += abstime.tv_sec + 
            (wakeuptime.tv_nsec + abstime.tv_usec * 1000) / 1000000000;
        wakeuptime.tv_nsec = 
            (wakeuptime.tv_nsec + abstime.tv_usec * 1000) % 1000000000;

#ifdef _DEBUG_PRINT
        printf("(Re)Transmitting seq %d, trials left %d, timeout %d...\n", 
                TCP_SEQ_NUM(node->packet), node->num_trials_left,
                timeout_msec);
#endif

        pthread_mutex_lock(&node->timer_kill_lock);
        if (node->timer_killed == TRUE)
            goto timer_killed;

        /* sleep and wait for kill signal */
        if (pthread_cond_timedwait(&node->timer_kill_cond,
                    &node->timer_kill_lock, &wakeuptime) == 0)
        {
            /* kill signal arrives, clean up and return */
            assert(node->timer_killed == TRUE);
            goto timer_killed;
        }
        else
        {
            /* timeout, retransmit packets.
             * Thought it is an error case, mutex is still acquired, so we have
             * to unlock it */

            if (node->timer_killed == TRUE)
                goto timer_killed;

            pthread_mutex_unlock(&node->timer_kill_lock);

            pthread_mutex_lock(&ctx->retransmission_lock);

            /* double the timeout */
            timeout_msec *= 2;

            /* only retransmit the first packet in queue */
            timer_queue_node_t *q = ctx->retransmission_queue->head;
            if (q)
                stcp_network_send(sd, q->packet, q->len, NULL);

            pthread_mutex_unlock(&ctx->retransmission_lock);

            node->num_trials_left--;
        }
    }

#ifdef _DEBUG_PRINT
    printf("[CONNECTION FAILED]\n");
#endif
    ctx->done = TRUE;
    ctx->connection_state = CLOSED;
    ctx->error = STCP_ERROR;

    return NULL;

timer_killed:
    pthread_mutex_unlock(&node->timer_kill_lock);
    free(arg);
    return NULL;
}

/* start a timer associated with the given node, using the timeout value in
 * the current context */
void start_timer(mysocket_t sd, context_t *ctx, timer_queue_node_t *node)
{
    struct timeval abstime; 
    
    /* We need to allocate memory for it. Puting it on the stack is not
     * enough, because the stack memory would be unaccessible after the
     * function call. The memory allocated here will be cleaned up in the
     * thread */
    timer_arg_type *arg = (timer_arg_type *)calloc(1, sizeof(timer_arg_type));

    arg->ctx = ctx;
    arg->sd = sd;
    arg->node = node;
    arg->timeout_msec = ctx->timeout;

    /* record current time */
    
    gettimeofday(&abstime, NULL);
    node->timer_start.tv_sec = abstime.tv_sec;
    node->timer_start.tv_nsec = abstime.tv_usec * 1000;

    /* start timer */
    pthread_create(&node->timer_thread, NULL, new_timer_thread_func, arg);
}

/* send packet and start a timer */
ssize_t send_and_start_timer(mysocket_t sd, context_t *ctx, const void *tcphdr,
        size_t hdr_len, const void *payload, size_t payload_len)
{
    timer_queue_t *rq;
    timer_queue_node_t *node;
    ssize_t rc;

    node = (timer_queue_node_t *) calloc(1, sizeof(timer_queue_node_t));
    node->packet = (uint8_t *) calloc(hdr_len + payload_len, sizeof(uint8_t));
    node->len = hdr_len + payload_len;
    memcpy(node->packet, tcphdr, hdr_len);
    if (payload_len > 0)
        memcpy(node->packet + hdr_len, payload, payload_len);

    /* timer_start is set later */
    node->num_trials_left = STCP_RETRANSMISSION_LIMIT; 
    node->next = NULL;
    node->timer_killed = FALSE;

    /* initialize condition variable and mutex */
    pthread_mutex_init(&node->timer_kill_lock, NULL);
    pthread_cond_init(&node->timer_kill_cond, NULL);

    /* send packet */
    rc = stcp_network_send(sd, node->packet, node->len, NULL);

    pthread_mutex_lock(&ctx->retransmission_lock);

    /* add the packet to retransmission queue */
    rq = ctx->retransmission_queue;
    if (!(rq->tail)) /* empty queue */
        rq->head = node;
    else
        rq->tail->next = node;
    rq->tail = node;

    pthread_mutex_unlock(&ctx->retransmission_lock);

    /* start timer */
    start_timer(sd, ctx, node);

    return rc;
}


/* send a pure flag packet, without any payload. Return -1 if there is a 
 * transmission error. The retransmission timer is started if the packet
 * contains flags other than ACK */
ssize_t send_flag_packet(mysocket_t sd, context_t *ctx, uint8_t flags)
{
    STCPHeader _tcphdr;
    STCPHeader *tcphdr = &_tcphdr;

    fill_in_tcphdr(tcphdr, ctx->send_next, ctx->recv_next, 
            STCP_DEFAULT_DATA_OFFSET, flags, ctx->recv_window_size);

#ifdef _DEBUG_PRINT
    print_packet_info(tcphdr, PACKET_SENT);
#endif

    if (flags & TH_FIN || flags & TH_SYN)
        return send_and_start_timer(sd, ctx, 
                tcphdr, sizeof(STCPHeader), NULL, 0);
    else
        return stcp_network_send(sd, tcphdr, sizeof(STCPHeader), NULL);
}

void print_tcphdr(STCPHeader *tcphdr)
{
    printf("----- STCP Header -----\n");
    printf("th_sport: %d\n", ntohs(tcphdr->th_sport));
    printf("th_dport: %d\n", ntohs(tcphdr->th_dport));
    printf("th_seq: %d\n", ntohl(tcphdr->th_seq));
    printf("th_ack: %d\n", ntohl(tcphdr->th_ack));
    printf("th_off: %d\n", tcphdr->th_off);
    printf("th_flags: %x\n", tcphdr->th_flags);
    printf("th_win: %d\n", ntohs(tcphdr->th_win));
    printf("th_sum: %d\n", tcphdr->th_sum);
    printf("th_urp: %d\n", tcphdr->th_urp);
    printf("\n");
}

void print_tcphdr_byte(uint8_t *tcphdr)
{
    uint8_t th_off = ((STCPHeader *)tcphdr)->th_off;
    int i = 0, j = 0;

    for (j = 0; j < th_off; j++)
    {
        for (i = 0; i < 4; i++)
            printf("%02X ", *(tcphdr + j*4 + i));
        printf("\n");
    }
}

/**********************************************************************/
/* our_dprintf
 *
 * Send a formatted message to stdout.
 * 
 * format               A printf-style format string.
 *
 * This function is equivalent to a printf, but may be
 * changed to log errors to a file if desired.
 *
 * Calls to this function are generated by the dprintf amd
 * dperror macros in transport.h
 */
void our_dprintf(const char *format,...)
{
    va_list argptr;
    char buffer[1024];

    assert(format);
    va_start(argptr, format);
    vsnprintf(buffer, sizeof(buffer), format, argptr);
    va_end(argptr);
    fputs(buffer, stdout);
    fflush(stdout);
}



